import json
import os
import tempfile

def process_jsonl_messages(jsonl_file_path):
    """
    处理JSONL文件中的消息记录：
    1. 持续删除最后一条消息，直到最后一条消息的role是assistant
    2. 如果删除后messages只剩下system消息，则删除整条记录
    """
    
    # 创建临时文件
    temp_file = tempfile.NamedTemporaryFile(mode='w', delete=False, encoding='utf-8')
    deleted_records = 0
    processed_records = 0
    
    try:
        with open(jsonl_file_path, 'r', encoding='utf-8') as file, temp_file:
            for line_num, line in enumerate(file, 1):
                line = line.strip()
                if not line:
                    continue
                    
                try:
                    data = json.loads(line)
                    messages = data.get('messages', [])
                    
                    # 检查是否有消息
                    if not messages:
                        # 如果消息列表为空，跳过此记录
                        continue
                    
                    # 持续删除最后一条消息，直到最后一条是assistant或消息列表为空
                    while messages and messages[-1].get('role') != 'assistant':
                        messages.pop()
                    
                    # 检查删除后是否只剩下system消息或为空
                    if not messages or all(msg.get('role') == 'system' for msg in messages):
                        # 如果只剩下system消息或为空，跳过此记录（不写入）
                        deleted_records += 1
                        continue
                    
                    # 更新数据并写入临时文件
                    data['messages'] = messages
                    json.dump(data, temp_file, ensure_ascii=False)
                    temp_file.write('\n')
                    processed_records += 1
                    
                except json.JSONDecodeError as e:
                    print(f"第{line_num}行JSON解析错误: {e}")
                    # 保留原行
                    temp_file.write(line + '\n')
                    processed_records += 1
                except Exception as e:
                    print(f"第{line_num}行处理错误: {e}")
                    # 保留原行
                    temp_file.write(line + '\n')
                    processed_records += 1
        
        # 用临时文件替换原文件
        temp_file_path = temp_file.name
        os.replace(temp_file_path, '为不同服务准备的数据集/OpenAI.jsonl')
        print(f"文件 {jsonl_file_path} 已成功处理")
        print(f"处理了 {processed_records} 条记录，删除了 {deleted_records} 条记录")
        
    except Exception as e:
        print(f"处理文件时发生错误: {e}")
        # 清理临时文件
        if 'temp_file_path' in locals():
            os.unlink(temp_file_path)

# 使用示例
if __name__ == "__main__":
    # 替换为你的JSONL文件路径
    file_path = "柏拉图.jsonl"
    process_jsonl_messages(file_path)